#include "peloader.h"
#include "kstdio.h"
#include "kmm.h"
#include "process.h"
#include "dbgio.h"

int exec_pe(char *filename){
	PIMAGE_DOS_HEADER dos_hdr;
	PIMAGE_NT_HEADERS nt_hdr;
	process_start_params_t * process_info;
	process_info_t *process;
	char fn[128];
	int file_len;
	int err;
	//open file
	strcpy(fn,filename);
	KFILE fd = open(filename,O_RDONLY,0);
	if (fd == -1){
		err = EIO;
		dprintf("ENOENT \n");
		goto loadpe_err_handler;
	}
	//Read MS-DOS headers
	dos_hdr = (PIMAGE_DOS_HEADER) kmalloc(sizeof(IMAGE_DOS_HEADER));
	if (read(fd,(uint8_t *) dos_hdr,sizeof(IMAGE_DOS_HEADER)) != sizeof(IMAGE_DOS_HEADER)){
		err = EIO;//Failed to read MZ header
		dprintf("ERMZ %i\n",get_current_process()->errno);
		goto loadpe_err_handler;
	}
	if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE){
		err = ENOMZ;//Not a MS-DOS/PE image
		dprintf("ENOMZ 0x%x\n",dos_hdr->e_magic);
		goto loadpe_err_handler;
	}
	if (dos_hdr->e_lfanew == 0){
		err = EMSDOS;//MS-DOS programs are not supported (yet)
		dprintf("EMSDOS 0x%x\n",dos_hdr->e_lfanew);
		goto loadpe_err_handler;
	}
	nt_hdr = (PIMAGE_NT_HEADERS) kmalloc(sizeof(IMAGE_NT_HEADERS));
	seek(fd,dos_hdr->e_lfanew,SEEK_SET);
	if (read(fd,(uint8_t *) nt_hdr,sizeof(IMAGE_NT_HEADERS)) != sizeof(IMAGE_NT_HEADERS)){
		err = EIO;//Failed to read NT headers
		dprintf("ERPE \n");
		goto loadpe_err_handler;
	}
	if (nt_hdr->Signature!= IMAGE_NT_SIGNATURE){
		err = ENOPE;//Not a PE image
		dprintf("ENOPE 0x%x\n",nt_hdr->Signature);
		goto loadpe_err_handler;
	}
	process_info = (process_start_params_t *) kmalloc(sizeof(process_start_params_t));
	process_info->entry_point = nt_hdr->OptionalHeader.AddressOfEntryPoint;
	process_info->initial_heap_size  = nt_hdr->OptionalHeader.SizeOfHeapCommit;
	process_info->initial_stack_size = nt_hdr->OptionalHeader.SizeOfStackCommit;
	process_info->max_heap_size = nt_hdr->OptionalHeader.SizeOfHeapReserve;
	process_info->max_stack_size = nt_hdr->OptionalHeader.SizeOfStackReserve;
	process_info->image_base = nt_hdr->OptionalHeader.ImageBase;
	file_len = lenf(fd);
	if (file_len == -1){
		err = EIO;
		dprintf("ERLEN \n");
		goto loadpe_err_handler;
	}
	process_info->code_size = file_len;
	process = create_process(process_info);
	seek(fd,0,SEEK_SET);
	if (read(fd,(uint8_t *)process->code_start,file_len) != file_len){
		err = EIO;
		dprintf("ERCODE \n");
		goto loadpe_err_handler;
	} 
	//dprintf("Starting task\n");
	strcpy(process->name,fn);
	return start_task(process);
loadpe_err_handler:
	if (dos_hdr != 0)
		kfree((void *) dos_hdr);
	if (nt_hdr != 0)
		kfree((void *) nt_hdr);
	if (process_info != 0)
		kfree((void *) process_info);
	if (fd != -1)
		close(fd);
	dprintf("Err %i \n",err);
	return err;

}